package com.github.binarylei.trident;

import org.apache.storm.Config;
import org.apache.storm.LocalCluster;
import org.apache.storm.StormSubmitter;
import org.apache.storm.generated.AlreadyAliveException;
import org.apache.storm.generated.AuthorizationException;
import org.apache.storm.generated.InvalidTopologyException;
import org.apache.storm.generated.StormTopology;
import org.apache.storm.trident.Stream;
import org.apache.storm.trident.TridentTopology;
import org.apache.storm.trident.operation.BaseFunction;
import org.apache.storm.trident.operation.TridentCollector;
import org.apache.storm.trident.testing.FixedBatchSpout;
import org.apache.storm.trident.tuple.TridentTuple;
import org.apache.storm.tuple.Fields;
import org.apache.storm.tuple.Values;

/**
 * @author: leigang
 * @version: 2018-04-10
 */
public class TridentFunction {

    public static void main(String[] args) throws InterruptedException, InvalidTopologyException,
            AuthorizationException, AlreadyAliveException {
        Config conf = new Config();
        conf.setNumWorkers(2);
        conf.setMaxSpoutPending(20);

        if (args == null || args.length ==0) {
            LocalCluster cluster = new LocalCluster();
            cluster.submitTopology("trident-function", conf, buildTopology());

            Thread.sleep(20 * 1000);
            cluster.shutdown();
        } else {
            StormSubmitter.submitTopology(args[0], conf, buildTopology());
        }
    }

    private static StormTopology buildTopology() {

        TridentTopology topology = new TridentTopology();
        // 设置数据源
        FixedBatchSpout spout = new FixedBatchSpout(
                new Fields("a", "b", "c", "d"),
                4,  // 设置批处理大小
                new Values(1, 4, 7, 10),
                new Values(1, 4, 7, 10),
                new Values(1, 4, 7, 10),
                new Values(1, 4, 7, 10)
        );
        // 指定是否循环
        spout.setCycle(false);
        Stream inputStream = topology.newStream("spout", spout);
        /**
         * 要实现流 spout-bout 的模式在 trident 里是使用 each 来做的
         * 1. 输入数据源参数名称："a", "b", "c", "d", "sum"
         * 2. 需要流转执行的 function 对象(也就是bolt对象)：new Result()
         * 3. 指定  function 对象里的输出参数名称：sum
         */
        inputStream.each(new Fields("a", "b", "c", "d"), new SumFunction(), new Fields("sum"))
                /**
                 * 继续使用 each 调用下一个 function(bolt)
                 * 第一个参数为："a", "b", "c", "d", "sum"
                 * 第二个参数为：new Result() 也就是执行函数
                 * 第三个参数为：new Fields() 没有输出
                 */
                .each(new Fields("a", "b", "c", "d", "sum"), new Result(), new Fields());
        return topology.build();
    }

    public static class SumFunction extends BaseFunction {

        @Override
        public void execute(TridentTuple tuple, TridentCollector collector) {
            System.out.println("传入进来的内容为：" + tuple);
            int a = tuple.getInteger(0);
            int b = tuple.getInteger(1);
            collector.emit(new Values(a + b));
        }
    }

    public static class Result extends BaseFunction {

        @Override
        public void execute(TridentTuple tuple, TridentCollector collector) {
            int a = tuple.getIntegerByField("a");
            int b = tuple.getIntegerByField("b");
            int c = tuple.getIntegerByField("c");
            int d = tuple.getIntegerByField("d");
            System.out.println(String.format("a: %s, b: %s, c: %s, d: %s", a, b, c, d));
            int sum = tuple.getIntegerByField("sum");
            System.out.println(String.format("sum:%s", sum));
        }
    }
}
